from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
import threading
import time
import concurrent.futures

msg = [
    SystemMessage("你是一位Python工程师，请回答以下问题"),
    HumanMessage("请问，Python的数据类型有哪些？")
]
total_time = 0


def task(name):
    global total_time
    print("start {} ...".format(name))
    start = time.time()

    parser = StrOutputParser()

    chain = model | parser

    result = chain.invoke(msg)

    print(result)
    end = time.time()
    total_time += (end - start)
    t = (end - start) / 60
    print("{} 花费了：{:.2f} 分钟".format(name, t))


class MyThread(threading.Thread):
    def __init__(self, name):
        super().__init__()
        self.name = name

    def run(self):  # 固定名字run ！！！必须用固定名
        task(self.name)


model = ChatOpenAI(model="deepseek-r1:14b",
                   openai_api_key="ollama",
                   openai_api_base="http://10.2.4.31:11434/v1/")

# model = ChatOpenAI(model="ernie-bot",
#                    openai_api_key="ollama",
#                    openai_api_base="http://10.2.4.161:8000/v1/")

# model = ChatOpenAI(model="deepseek-r1-qwen-32b",
#                    openai_api_key="ollama",
#                    openai_api_base="http://10.2.42.9:8112/v1")

if __name__ == '__main__':  # 必须要这样启动
    threads = []
    for i in range(1):
        p = MyThread('聊天任务' + str(i))
        threads.append(p)
        p.start()
    # 等待所有线程任务结束
    for t in threads:
        t.join()
    print("总花费时间：{:.2f} 分钟".format(total_time / 60))
    print("平均时长：{:.2f}".format(total_time / 40))
